/*
 * -----------------------------------------------------------------
 * Programmer(s): Scott D. Cohen, Alan C. Hindmarsh, George Byrne,
 *                and Radu Serban @ LLNL
 * -----------------------------------------------------------------
 * SUNDIALS Copyright Start
 * Copyright (c) 2002-2025, Lawrence Livermore National Security
 * and Southern Methodist University.
 * All rights reserved.
 *
 * See the top-level LICENSE and NOTICE files for details.
 *
 * SPDX-License-Identifier: BSD-3-Clause
 * SUNDIALS Copyright End
 * -----------------------------------------------------------------
 * Example problem:
 *
 * The following is a simple example problem, with the program for
 * its solution by CVODE. The problem is the semi-discrete
 * form of the advection-diffusion equation in 1-D:
 *   du/dt = d^2 u / dx^2 + .5 du/dx
 * on the interval 0 <= x <= 2, and the time interval 0 <= t <= 5.
 * Homogeneous Dirichlet boundary conditions are posed, and the
 * initial condition is the following:
 *   u(x,t=0) = x(2-x)exp(2x) .
 * The PDE is discretized on a uniform grid of size MX+2 with
 * central differencing, and with boundary values eliminated,
 * leaving an ODE system of size NEQ = MX.
 * This program solves the problem with the ADAMS integration method,
 * and with Newton iteration using diagonal approximate Jacobians.
 * It uses scalar relative and absolute tolerances.
 * Output is printed at t = .5, 1.0, ..., 5.
 * Run statistics (optional outputs) are printed at the end.
 *
 * This version uses MPI for user routines.
 * Execute with Number of Processors = N,  with 1 <= N <= MX.
 * -----------------------------------------------------------------
 */

#include <cvode/cvode.h>      /* prototypes for CVODE fcts., consts. */
#include <cvode/cvode_diag.h> /* prototypes for CVODE diagonal solver */
#include <math.h>
#include <mpi.h>                      /* MPI constants and types */
#include <nvector/nvector_parallel.h> /* access to MPI-parallel N_Vector     */
#include <stdio.h>
#include <stdlib.h>
#include <sundials/sundials_core.h>

/* Problem Constants */

#define ZERO SUN_RCONST(0.0)

#define XMAX  SUN_RCONST(2.0)    /* domain boundary           */
#define MX    10                 /* mesh dimension            */
#define NEQ   MX                 /* number of equations       */
#define ATOL  SUN_RCONST(1.0e-5) /* scalar absolute tolerance */
#define T0    ZERO               /* initial time              */
#define T1    SUN_RCONST(0.5)    /* first output time         */
#define DTOUT SUN_RCONST(0.5)    /* output time increment     */
#define NOUT  10                 /* number of output times    */

/* Type : UserData
   contains grid constants, parallel machine parameters, work array. */

typedef struct
{
  sunrealtype dx, hdcoef, hacoef;
  int npes, my_pe;
  MPI_Comm comm;
  sunrealtype z[100];
}* UserData;

/* Private Helper Functions */

static void SetIC(N_Vector u, sunrealtype dx, sunindextype my_length,
                  sunindextype my_base);

static void PrintIntro(int npes);

static void PrintData(sunrealtype t, sunrealtype umax, long int nst);

static void PrintFinalStats(void* cvode_mem);

/* Functions Called by the Solver */

static int f(sunrealtype t, N_Vector u, N_Vector udot, void* user_data);

/* Private function to check function return values */

static int check_retval(void* returnvalue, const char* funcname, int opt, int id);

/***************************** Main Program ******************************/

int main(int argc, char* argv[])
{
  SUNContext sunctx;
  SUNLogger logger;
  sunrealtype dx, reltol, abstol, t, tout, umax;
  N_Vector u;
  UserData data;
  void* cvode_mem;
  int iout, retval, my_pe, npes;
  sunindextype local_N, nperpe, nrem, my_base;
  long int nst;

  MPI_Comm comm;

  sunctx    = NULL;
  logger    = NULL;
  u         = NULL;
  data      = NULL;
  cvode_mem = NULL;

  /* Get processor number, total number of pe's, and my_pe. */
  MPI_Init(&argc, &argv);
  comm = MPI_COMM_WORLD;
  MPI_Comm_size(comm, &npes);
  MPI_Comm_rank(comm, &my_pe);

  /* Create the SUNDIALS context */
  retval = SUNContext_Create(comm, &sunctx);
  if (check_retval(&retval, "SUNContext_Create", 1, my_pe))
  {
    MPI_Abort(comm, 1);
  }

  /* Setup different error handler stack so that we abort after logging */
  SUNContext_PopErrHandler(sunctx);
  SUNContext_PushErrHandler(sunctx, SUNMPIAbortErrHandlerFn, NULL);
  SUNContext_PushErrHandler(sunctx, SUNLogErrHandlerFn, NULL);

  /* This requires that SUNDIALS was configured with the CMake options
       SUNDIALS_LOGGING_LEVEL=n
    where n is one of:
       1 --> log only errors,
       2 --> log errors + warnings,
       3 --> log errors + warnings + info output
       4 --> all of the above plus debugging output like internal integrator values
       5 --> all of the above and even more
    SUNDIALS will only log up to the max level n, but a lesser level can
    be configured at runtime by only providing output files for the
    desired levels. We will enable all logging here on the condition
    that SUNDIALS was built with the correct logging level enabled. */

  if (SUNDIALS_LOGGING_LEVEL >= SUN_LOGLEVEL_ERROR)
  {
    retval = SUNLogger_Create(comm, // MPI communicator
                              0,    // output on process 0
                              &logger);
    if (check_retval(&retval, "SUNLogger_Create", 1, my_pe))
    {
      MPI_Abort(comm, 1);
    }

    retval = SUNContext_SetLogger(sunctx, logger);
    if (check_retval(&retval, "SUNContext_SetLogger", 1, my_pe))
    {
      MPI_Abort(comm, 1);
    }

    retval = SUNLogger_SetErrorFilename(logger, "stderr");
    if (check_retval(&retval, "SUNLogger_SetErrorFilename", 1, my_pe))
    {
      MPI_Abort(comm, 1);
    }
  }

  if (SUNDIALS_LOGGING_LEVEL >= SUN_LOGLEVEL_WARNING)
  {
    retval = SUNLogger_SetWarningFilename(logger, "stderr");
    if (check_retval(&retval, "SUNLogger_SetWarningFilename", 1, my_pe))
    {
      MPI_Abort(comm, 1);
    }
  }

  if (SUNDIALS_LOGGING_LEVEL >= SUN_LOGLEVEL_INFO)
  {
    retval = SUNLogger_SetInfoFilename(logger, "cvAdvDiff_diag_p.info.log");
    if (check_retval(&retval, "SUNLogger_SetInfoFilename", 1, my_pe))
    {
      MPI_Abort(comm, 1);
    }
  }

  if (SUNDIALS_LOGGING_LEVEL >= SUN_LOGLEVEL_DEBUG)
  {
    retval = SUNLogger_SetDebugFilename(logger, "stderr");
    if (check_retval(&retval, "SUNLogger_SetDebugFilename", 1, my_pe))
    {
      MPI_Abort(comm, 1);
    }
  }

  /* Set local vector length. */
  nperpe  = NEQ / npes;
  nrem    = NEQ - npes * nperpe;
  local_N = (my_pe < nrem) ? nperpe + 1 : nperpe;
  my_base = (my_pe < nrem) ? my_pe * local_N : my_pe * nperpe + nrem;

  data = (UserData)malloc(sizeof *data); /* Allocate data memory */
  if (check_retval((void*)data, "malloc", 2, my_pe)) { MPI_Abort(comm, 1); }

  data->comm  = comm;
  data->npes  = npes;
  data->my_pe = my_pe;

  u = N_VNew_Parallel(comm, local_N, NEQ, sunctx); /* Allocate u vector */
  if (check_retval((void*)u, "N_VNew", 0, my_pe)) { MPI_Abort(comm, 1); }

  reltol = ZERO; /* Set the tolerances */
  abstol = ATOL;

  dx = data->dx = XMAX /
                  ((sunrealtype)(MX + 1)); /* Set grid coefficients in data */
  data->hdcoef = SUN_RCONST(1.0) / (dx * dx);
  data->hacoef = SUN_RCONST(0.5) / (SUN_RCONST(2.0) * dx);

  SetIC(u, dx, local_N, my_base); /* Initialize u vector */

  /* Call CVodeCreate to create the solver memory and specify the
   * Adams-Moulton LMM */
  cvode_mem = CVodeCreate(CV_ADAMS, sunctx);
  if (check_retval((void*)cvode_mem, "CVodeCreate", 0, my_pe))
  {
    MPI_Abort(comm, 1);
  }

  retval = CVodeSetUserData(cvode_mem, data);
  if (check_retval(&retval, "CVodeSetUserData", 1, my_pe))
  {
    MPI_Abort(comm, 1);
  }

  /* Call CVodeInit to initialize the integrator memory and specify the
   * user's right hand side function in u'=f(t,u), the initial time T0, and
   * the initial dependent variable vector u. */
  retval = CVodeInit(cvode_mem, f, T0, u);
  if (check_retval(&retval, "CVodeInit", 1, my_pe)) { return (1); }

  /* Call CVodeSStolerances to specify the scalar relative tolerance
   * and scalar absolute tolerances */
  retval = CVodeSStolerances(cvode_mem, reltol, abstol);
  if (check_retval(&retval, "CVodeSStolerances", 1, my_pe)) { return (1); }

  /* Call CVDiag to create and attach CVODE-specific diagonal linear solver */
  retval = CVDiag(cvode_mem);
  if (check_retval(&retval, "CVDiag", 1, my_pe)) { return (1); }

  if (my_pe == 0) { PrintIntro(npes); }

  umax = N_VMaxNorm(u);

  if (my_pe == 0)
  {
    t = T0;
    PrintData(t, umax, 0);
  }

  /* In loop over output points, call CVode, print results, test for error */

  for (iout = 1, tout = T1; iout <= NOUT; iout++, tout += DTOUT)
  {
    retval = CVode(cvode_mem, tout, u, &t, CV_NORMAL);
    if (check_retval(&retval, "CVode", 1, my_pe)) { break; }
    umax   = N_VMaxNorm(u);
    retval = CVodeGetNumSteps(cvode_mem, &nst);
    check_retval(&retval, "CVodeGetNumSteps", 1, my_pe);
    if (my_pe == 0) { PrintData(t, umax, nst); }
  }

  if (my_pe == 0)
  {
    PrintFinalStats(cvode_mem); /* Print some final statistics */
  }

  N_VDestroy(u);         /* Free the u vector */
  CVodeFree(&cvode_mem); /* Free the integrator memory */
  free(data);            /* Free user data */
  if (logger) { SUNLogger_Destroy(&logger); }
  SUNContext_Free(&sunctx);

  MPI_Finalize();
  return (0);
}

/************************ Private Helper Functions ***********************/

/* Set initial conditions in u vector */

static void SetIC(N_Vector u, sunrealtype dx, sunindextype my_length,
                  sunindextype my_base)
{
  int i;
  sunindextype iglobal;
  sunrealtype x;
  sunrealtype* udata;

  /* Set pointer to data array and get local length of u. */
  udata     = N_VGetArrayPointer(u);
  my_length = N_VGetLocalLength_Parallel(u);

  /* Load initial profile into u vector */
  for (i = 1; i <= my_length; i++)
  {
    iglobal      = my_base + i;
    x            = iglobal * dx;
    udata[i - 1] = x * (XMAX - x) * exp(SUN_RCONST(2.0) * x);
  }
}

/* Print problem introduction */

static void PrintIntro(int npes)
{
  printf("\n 1-D advection-diffusion equation, mesh size =%3d \n", MX);
  printf("\n Number of PEs = %3d \n", npes);
  printf("\n Diagonal linear solver CVDiag \n\n");

  return;
}

/* Print data */

static void PrintData(sunrealtype t, sunrealtype umax, long int nst)
{
#if defined(SUNDIALS_EXTENDED_PRECISION)
  printf("At t = %4.2Lf  max.norm(u) =%14.6Le  nst =%4ld \n", t, umax, nst);
#elif defined(SUNDIALS_DOUBLE_PRECISION)
  printf("At t = %4.2f  max.norm(u) =%14.6e  nst =%4ld \n", t, umax, nst);
#else
  printf("At t = %4.2f  max.norm(u) =%14.6e  nst =%4ld \n", t, umax, nst);
#endif

  return;
}

/* Print some final statistics located in the iopt array */

static void PrintFinalStats(void* cvode_mem)
{
  long int nst, nfe, nni, ncfn, netf;
  int retval;

  retval = CVodeGetNumSteps(cvode_mem, &nst);
  check_retval(&retval, "CVodeGetNumSteps", 1, 0);
  retval = CVodeGetNumRhsEvals(cvode_mem, &nfe);
  check_retval(&retval, "CVodeGetNumRhsEvals", 1, 0);
  retval = CVodeGetNumErrTestFails(cvode_mem, &netf);
  check_retval(&retval, "CVodeGetNumErrTestFails", 1, 0);
  retval = CVodeGetNumNonlinSolvIters(cvode_mem, &nni);
  check_retval(&retval, "CVodeGetNumNonlinSolvIters", 1, 0);
  retval = CVodeGetNumNonlinSolvConvFails(cvode_mem, &ncfn);
  check_retval(&retval, "CVodeGetNumNonlinSolvConvFails", 1, 0);

  printf("\nFinal Statistics: \n\n");
  printf("nst = %-6ld  nfe  = %-6ld  ", nst, nfe);
  printf("nni = %-6ld  ncfn = %-6ld  netf = %ld\n \n", nni, ncfn, netf);
}

/***************** Function Called by the Solver ***********************/

/* f routine. Compute f(t,u). */

static int f(sunrealtype t, N_Vector u, N_Vector udot, void* user_data)
{
  sunrealtype ui, ult, urt, hordc, horac, hdiff, hadv;
  sunrealtype *udata, *dudata, *z;
  int npes, my_pe, my_pe_m1, my_pe_p1, last_pe;
  sunindextype i, my_length;
  UserData data;
  MPI_Status status;
  MPI_Comm comm;

  udata  = N_VGetArrayPointer(u);
  dudata = N_VGetArrayPointer(udot);

  /* Extract needed problem constants from data */
  data  = (UserData)user_data;
  hordc = data->hdcoef;
  horac = data->hacoef;

  /* Extract parameters for parallel computation. */
  comm  = data->comm;
  npes  = data->npes;                        /* Number of processes. */
  my_pe = data->my_pe;                       /* Current process number. */
  my_length = N_VGetLocalLength_Parallel(u); /* Number of local elements of u. */
  z = data->z;

  /* Compute related parameters. */
  my_pe_m1 = my_pe - 1;
  my_pe_p1 = my_pe + 1;
  last_pe  = npes - 1;

  /* Store local segment of u in the working array z. */
  for (i = 1; i <= my_length; i++) { z[i] = udata[i - 1]; }

  /* Pass needed data to processes before and after current process. */
  if (my_pe != 0) { MPI_Send(&z[1], 1, MPI_SUNREALTYPE, my_pe_m1, 0, comm); }
  if (my_pe != last_pe)
  {
    MPI_Send(&z[my_length], 1, MPI_SUNREALTYPE, my_pe_p1, 0, comm);
  }

  /* Receive needed data from processes before and after current process. */
  if (my_pe != 0)
  {
    MPI_Recv(&z[0], 1, MPI_SUNREALTYPE, my_pe_m1, 0, comm, &status);
  }
  else { z[0] = ZERO; }
  if (my_pe != last_pe)
  {
    MPI_Recv(&z[my_length + 1], 1, MPI_SUNREALTYPE, my_pe_p1, 0, comm, &status);
  }
  else { z[my_length + 1] = ZERO; }

  /* Loop over all grid points in current process. */
  for (i = 1; i <= my_length; i++)
  {
    /* Extract u at x_i and two neighboring points */
    ui  = z[i];
    ult = z[i - 1];
    urt = z[i + 1];

    /* Set diffusion and advection terms and load into udot */
    hdiff         = hordc * (ult - SUN_RCONST(2.0) * ui + urt);
    hadv          = horac * (urt - ult);
    dudata[i - 1] = hdiff + hadv;
  }

  return (0);
}

/* Check function return value...
     opt == 0 means SUNDIALS function allocates memory so check if
              returned NULL pointer
     opt == 1 means SUNDIALS function returns an integer value so check if
              retval < 0
     opt == 2 means function allocates memory so check if returned
              NULL pointer */

static int check_retval(void* returnvalue, const char* funcname, int opt, int id)
{
  int* retval;

  /* Check if SUNDIALS function returned NULL pointer - no memory allocated */
  if (opt == 0 && returnvalue == NULL)
  {
    fprintf(stderr,
            "\nSUNDIALS_ERROR(%d): %s() failed - returned NULL pointer\n\n", id,
            funcname);
    return (1);
  }

  /* Check if retval < 0 */
  else if (opt == 1)
  {
    retval = (int*)returnvalue;
    if (*retval < 0)
    {
      fprintf(stderr, "\nSUNDIALS_ERROR(%d): %s() failed with retval = %d\n\n",
              id, funcname, *retval);
      return (1);
    }
  }

  /* Check if function returned NULL pointer - no memory allocated */
  else if (opt == 2 && returnvalue == NULL)
  {
    fprintf(stderr, "\nMEMORY_ERROR(%d): %s() failed - returned NULL pointer\n\n",
            id, funcname);
    return (1);
  }

  return (0);
}
